from calc_manager.plugins_inner.bdf_reader_inner import BDFReader
from calc_manager.plugins_inner.op2_reader_inner import OP2Reader
from calc_manager.plugins_inner.mat_csv_inner import MatCSVReader
from calc_manager.plugins_inner.bdf_transformer_inner import BDFTransformer
from pydantic import BaseModel
from typing import List, Dict


class Parameter(BaseModel):
    margin_output: List[Dict]
    danger_infos: Dict[str, List[float]]
    rough_eids: List[float]

class Result(BaseModel):
    stress_output: List[Dict]


def third(param: Parameter, bdf_transformer: BDFTransformer,op2_reader: OP2Reader) :
    print("step3-----------------> start")
    """
    第三张云图,只需要g70_name的数据
    1. gname代表表格第一列的G70编号结果
    2. 基于eids和stress数据绘制第三张应力云图
    """
    # TODO 这个数据可以和第二张图的数据合并
    danger_infos = param.danger_infos
    rough_eids = param.rough_eids
    rough_eids_set = set(rough_eids) # Convert to set for O(1) lookup


    results = {}
    subcase_id = op2_reader.get_subcase_id()
    result_components = ["stress.ctria3_stress", "stress.cquad4_stress"]
    for comp in result_components:
        print(f"step3-----------------> Getting results for component {comp} and subcase {subcase_id}")
        stress = op2_reader.get_result_from_comp_and_subcase(comp, subcase_id)
        print(f"step3-----------------> Got stress results with shape: {stress.data.shape}")
        element_node = stress.element_node
        von_mises = stress.von_mises()
        for index, item in enumerate(element_node):
            if index % 2 == 0:
                eid = item[0]
                # 只保存非rough_eids的结果 , 过滤掉自然网格
                if eid not in rough_eids_set: # 使用set进行O(1)的查找
                    results[eid] = von_mises[0][index : index + 2].tolist()
    print(f"step3-----------------> results: complete") 

    #results = {}
    #subcase_id = op2_reader.get_subcase_id()
    #result_components = ["stress.ctria3_stress", "stress.cquad4_stress"]
    #for comp in result_components:
    #    print(f"Getting results for component {comp} and subcase {subcase_id}")
    #    stress = op2_reader.get_result_from_comp_and_subcase(comp, subcase_id)
    #    print(f"Got stress results with shape: {stress.data.shape}")
    #    element_node = stress.element_node
    #    von_mises = stress.von_mises()
    #    for index, item in enumerate(element_node):
    #        if index % 2 == 0:
    #            results[item[0]] = von_mises[0][index : index + 2].tolist()

    ## 删除自然网格的结果(无需关注)
    #for eid in rough_eids:
    #    results.pop(eid, None)


    stress_output = []
    for k, v in danger_infos.items():
        eids = bdf_transformer.get_comps_by_gname(k)[1]
        stress = [max(results[eid]) for eid in eids]
        stress_output.append(
            {
                "gname": k,
                "eids": eids,
                "stress": stress
            }
        )

    #### 如果通过算子算法平台计算,则需要返回以下数据
    """
    1. rough_eids, refine_eids
    2. margin_output
    3. stress_output
    """

    res =  Result(stress_output=stress_output)
    print("step3-----------------> end")
    return res
